'''
Created on Jan 29, 2012

@author: nzhao
'''
import math
from const import *

class NVCenter:
    """Class of a single NV center"""
    def __init__(self, b = 0.0010, theta = 0.0, phi = 0.0, nuc='None'):
        self.b = b
        self.theta = theta / 180. * math.pi 
        self.phi = phi / 180. * math.pi
        self.nuc = nuc
        
        if nuc == 'N14':
            self.nuc_dim = 3
            self.IX, self.IY, self.IZ = SPIN1X, SPIN1Y, SPIN1Z
        elif nuc == 'N15':
            self.nuc_dim = 2
            self.IX, self.IY, self.IZ = 0.5 * X, 0.5 * Y, 0.5 * Z
        else:
            self.nuc_dim = 1
            self.IX, self.IY, self.IZ = 1, 1, 1
        
        self.SXop = np.kron(SPIN1X, np.identity(self.nuc_dim))
        self.SYop = np.kron(SPIN1Y, np.identity(self.nuc_dim))
        self.SZop = np.kron(SPIN1Z, np.identity(self.nuc_dim))

        self.IXop = np.kron(np.identity(3), self.IX)
        self.IYop = np.kron(np.identity(3), self.IY)
        self.IZop = np.kron(np.identity(3), self.IZ)
        
        self.calc_hamiltonian()
        
    def calc_hamiltonian(self):
        self.nuc_hamiltonian = self.hami_nuc_spin()
        self.e_hamiltonian = self.hami_e_spin()
        self.hamiltonian = self.hami()
        
    def bList(self):
        bx = self.b * math.sin(self.theta) * math.cos(self.phi)
        by = self.b * math.sin(self.theta) * math.sin(self.phi)
        bz = self.b * math.cos(self.theta)
        return np.array([bx, by, bz])
    
    def hami_nuc_spin(self):
        blist = self.bList()
        if self.nuc == 'N14':
            [bx, by, bz] = - Gamma_N14 * blist / 1000. # k rad / s
            res = bx * SPIN1X + by * SPIN1Y + bz * SPIN1Z \
                + 2.0 * math.pi * NV_N14_ZFS * 10**3 * np.dot(SPIN1Z, SPIN1Z)
        elif self.nuc == 'N15':
            [bx, by, bz] = - Gamma_N15 * blist / 1000. # k rad / s
            res = bx * X + by * Y + bz * Z
        else:
            res = 1.;
        return res
    
    def hami_e_spin(self):
        [bx, by, bz] = - GAMMA_E * self.bList() / 1000. # k rad / s 
        res = bx * SPIN1X + by * SPIN1Y + bz * SPIN1Z \
            + 2.0 * math.pi * NV_ZFS * 10**6 * np.dot(SPIN1Z, SPIN1Z)
        return res
    
    def hami(self):
        res = np.kron(self.e_hamiltonian, self.nuc_hamiltonian)
        return res
    
    def eigensys(self):
        nlen = len(self.hamiltonian)
        
        eig_freq, eig_wavefunc = np.linalg.eig(self.hamiltonian)
        sort_idx = eig_freq.argsort()
        eig_freq = eig_freq[sort_idx]
        eig_wavefunc = eig_wavefunc[:,sort_idx]
        
        statelist = [{'freq': eig_freq[i], 'wave_func': eig_wavefunc[i]} 
                     for i in range(nlen)]
        
        self.eigen_states = map(self.state_calc, statelist)
        
    def state_calc(self, state):
        wf = state['wave_func']
        wfc = wf.conj()
        
        Svec =  map(lambda op: np.dot(wfc, np.dot(op, wf)), 
                    [self.SXop, self.SYop, self.SZop])
        Ivec =  map(lambda op: np.dot(wfc, np.dot(op, wf)), 
                    [self.IXop, self.IYop, self.IZop])
        state['S_vector'] = np.array(Svec)
        state['I_vector'] = np.array(Ivec)
        res = state
        return res        